const fs = require('fs');
const path = require('path');

/**
 * Processes a string (typically from an AI response) to extract code blocks,
 * save them to files, and optionally create an explanation file.
 *
 * @param {string} aiResponse The string content to process.
 * @param {string} finalBaseName The base name to use for output files.
 * @param {string} inputDir The directory where output files will be saved.
 * @param {string|null} originalInputPath The path of the original input file, used to prevent overwriting.
 * @param {string} inputExt The extension of the input file, used as a fallback for single code blocks.
 * @returns {{originalLines: number, filesCreated: number, totalExtractedLines: number, remainingLines: number, filesCreatedDetails: Array<{path: string, lines: number}>}} An object summarizing the processing results.
 */
function processAIResponse(aiResponse, finalBaseName, inputDir, originalInputPath = null, inputExt = '.cu') {
  const originalLines = aiResponse.split('\n').length;

  // Extract code blocks with specific program names or numbered patterns
  const specificHeaderRegex = /\*\*\d+\.\s*`([^`]+\.cuh)`\*\*\s*```[a-zA-Z]*\r?\n([\s\S]*?)\r?\n```/;
  const specificImplementationRegex = /\*\*\d+\.\s*`([^`]+\.cu)`\*\*\s*```[a-zA-Z]*\r?\n([\s\S]*?)\r?\n```/;
  
  // Additional patterns for handling cases without markdown formatting
  const simpleHeaderRegex = /```cuda\r?\n(#ifndef[\s\S]*?#endif)\r?\n```/;
  const simpleImplementationRegex = /```cuda\r?\n(#include[\s\S]*?(?:__global__|void\s+\w+)[\s\S]*?)\r?\n```/;

  // Also try generic patterns without specific filenames
  const genericHeaderRegex = /\*\*1\.\s*(?!`)([^*]*?)\*\*\s*```[a-zA-Z]*\r?\n([\s\S]*?)\r?\n```/;
  const genericImplementationRegex = /\*\*2\.\s*(?!`)([^*]*?)\*\*\s*```[a-zA-Z]*\r?\n([\s\S]*?)\r?\n```/;

  let headerMatch = aiResponse.match(specificHeaderRegex);
  let implementationMatch = aiResponse.match(specificImplementationRegex);

  let headerFileName = null;
  let implementationFileName = null;

  if (headerMatch && headerMatch[1]) {
    headerFileName = path.basename(headerMatch[1], '.cuh');
  } else {
    headerMatch = aiResponse.match(genericHeaderRegex);
    if (!headerMatch) {
      // Try simple header pattern
      const simpleHeader = aiResponse.match(simpleHeaderRegex);
      if (simpleHeader) {
        headerMatch = [simpleHeader[0], 'vector_add.cuh', simpleHeader[1]];
        headerFileName = 'vector_add';
      }
    }
  }

  if (implementationMatch && implementationMatch[1]) {
    implementationFileName = path.basename(implementationMatch[1], '.cu');
  } else {
    implementationMatch = aiResponse.match(genericImplementationRegex);
    if (!implementationMatch) {
      // Try simple implementation pattern
      const simpleImpl = aiResponse.match(simpleImplementationRegex);
      if (simpleImpl) {
        implementationMatch = [simpleImpl[0], 'vector_add.cu', simpleImpl[1]];
        implementationFileName = 'vector_add';
      }
    }
  }

  let singleCodeBlockMatch = null;
  if (!headerMatch && !implementationMatch) {
    const singleCodeBlockRegex = /```[a-zA-Z]*\r?\n([\s\S]*?)\r?\n```/;
    singleCodeBlockMatch = aiResponse.match(singleCodeBlockRegex);
  }

  let filesCreated = 0;
  let totalExtractedLines = 0;
  let filesCreatedDetails = [];

  if (headerMatch && headerMatch[2]) {
    const headerContent = headerMatch[2];
    const headerLines = headerContent.split('\n').length;
    
    let actualHeaderFile;
    if (headerFileName) {
      actualHeaderFile = path.join(inputDir, headerFileName + '.cuh');
    } else {
      actualHeaderFile = path.join(inputDir, finalBaseName + '_part1.cuh');
    }
    
    fs.writeFileSync(actualHeaderFile, headerContent, 'utf8');
    console.log(`✓ Header file: ${path.basename(actualHeaderFile)} (${headerLines} lines)`);
    filesCreated++;
    totalExtractedLines += headerLines;
    filesCreatedDetails.push({ path: actualHeaderFile, lines: headerLines });
  } else {
    console.log('No header code block found.');
  }

  if (implementationMatch && implementationMatch[2]) {
    const implementationContent = implementationMatch[2];
    const implementationLines = implementationContent.split('\n').length;
    
    let actualImplementationFile;
    if (implementationFileName) {
      actualImplementationFile = path.join(inputDir, implementationFileName + '.cu');
    } else {
      actualImplementationFile = path.join(inputDir, finalBaseName + '_part2.cu');
    }
    
    fs.writeFileSync(actualImplementationFile, implementationContent, 'utf8');
    console.log(`✓ Implementation file: ${path.basename(actualImplementationFile)} (${implementationLines} lines)`);
    filesCreated++;
    totalExtractedLines += implementationLines;
    filesCreatedDetails.push({ path: actualImplementationFile, lines: implementationLines });
  } else {
    console.log('No implementation code block found.');
  }

  if (singleCodeBlockMatch && singleCodeBlockMatch[1]) {
    const codeContent = singleCodeBlockMatch[1];
    const codeLines = codeContent.split('\n').length;
    
    let singleOutputFile;
    if (inputExt) {
      singleOutputFile = path.join(inputDir, finalBaseName + '_extracted' + inputExt);
    } else {
      singleOutputFile = path.join(inputDir, finalBaseName + '_extracted.cu');
    }
    
    fs.writeFileSync(singleOutputFile, codeContent, 'utf8');
    console.log(`✓ Code file: ${path.basename(singleOutputFile)} (${codeLines} lines)`);
    filesCreated++;
    totalExtractedLines += codeLines;
    filesCreatedDetails.push({ path: singleOutputFile, lines: codeLines });
  } else if (!headerMatch && !implementationMatch) {
    const codePatterns = [
      /^#ifndef\s+\w+/m, /^#include\s+[<"]/m, /^\/\*\*?/m, /^\/\//m,
      /^class\s+\w+/m, /^struct\s+\w+/m, /^namespace\s+\w+/m,
      /^__global__/m, /^__device__/m, /^template\s*</m
    ];
    
    const isCodeFile = codePatterns.some(pattern => pattern.test(aiResponse));
    
    if (isCodeFile) {
      let rawCodeOutputFile;
      if (inputExt) {
        rawCodeOutputFile = path.join(inputDir, finalBaseName + '_extracted' + inputExt);
      } else {
        rawCodeOutputFile = path.join(inputDir, finalBaseName + '_extracted.cu');
      }
      
      const resolvedRawCodePath = path.resolve(rawCodeOutputFile);
      const resolvedOriginalInputPath = originalInputPath ? path.resolve(originalInputPath) : null;

      if (resolvedRawCodePath !== resolvedOriginalInputPath) {
        fs.writeFileSync(rawCodeOutputFile, aiResponse, 'utf8');
        console.log(`✓ Raw code file: ${path.basename(rawCodeOutputFile)} (${originalLines} lines)`);
        filesCreated++;
        totalExtractedLines += originalLines;
        filesCreatedDetails.push({ path: rawCodeOutputFile, lines: originalLines });
      } else {
        console.log(`Raw code file detected: ${path.basename(originalInputPath)} (${originalLines} lines) - no extraction needed`);
      }
    }
  }

  // Count total codeblocks found
  let totalCodeblocks = 0;
  if (headerMatch) totalCodeblocks++;
  if (implementationMatch) totalCodeblocks++;
  if (singleCodeBlockMatch) totalCodeblocks++;

  let remainingContent = aiResponse;
  let remainingLines = 0;

  // Create llm_output directory and save original LLM response
  const llmOutDir = path.join(inputDir, 'llm_output');
  if (!fs.existsSync(llmOutDir)) {
    fs.mkdirSync(llmOutDir, { recursive: true });
  }
  
  const llmOutputFile = path.join(llmOutDir, finalBaseName + '.cu');
  fs.writeFileSync(llmOutputFile, aiResponse, 'utf8');
  console.log(`✓ Original LLM response saved to: llm_output/${path.basename(llmOutputFile)}`);

  if (filesCreated > 0) {
    if (headerMatch && headerMatch[0]) {
      remainingContent = remainingContent.replace(new RegExp(headerMatch[0].replace(/[.*+?^${}()|[\]\\]/g, '\\$&'), 'g'), '');
    }
    if (implementationMatch && implementationMatch[0]) {
      remainingContent = remainingContent.replace(new RegExp(implementationMatch[0].replace(/[.*+?^${}()|[\]\\]/g, '\\$&'), 'g'), '');
    }
    if (singleCodeBlockMatch && singleCodeBlockMatch[0]) {
      remainingContent = remainingContent.replace(new RegExp(singleCodeBlockMatch[0].replace(/[.*+?^${}()|[\]\\]/g, '\\$&'), 'g'), '');
    }
    
    remainingContent = remainingContent.trim();
    remainingLines = remainingContent ? remainingContent.split('\n').length : 0;
    
    if (remainingLines > 10) {
      const explanationFile = path.join(inputDir, finalBaseName + '_explanation.md');
      fs.writeFileSync(explanationFile, remainingContent, 'utf8');
      console.log(`✓ Explanation file: ${path.basename(explanationFile)} (${remainingLines} lines)`);
      filesCreated++;
      filesCreatedDetails.push({ path: explanationFile, lines: remainingLines });
    }
  } else {
    remainingContent = remainingContent.trim();
    remainingLines = remainingContent ? remainingContent.split('\n').length : 0;
  }

  // Print summary like simulate_processing.js
  if (filesCreated > 0) {
    console.log('\n--- Processing Summary ---');
    console.log(`Original input: ${originalLines} lines`);
    console.log(`Codeblocks identified: ${totalCodeblocks}`);
    console.log(`Files created: ${filesCreated}`);
    if (totalExtractedLines > 0) {
      console.log(`Total extracted code lines: ${totalExtractedLines}`);
    }
    console.log(`Remaining content: ${remainingLines} lines`);
    
    console.log('\n--- Created Files ---');
    filesCreatedDetails.forEach(file => {
      const fileType = path.extname(file.path) === '.md' ? 'Documentation' : 'Code';
      console.log(`${fileType}: ${path.basename(file.path)} (${file.lines} lines)`);
    });
  }

  return {
    originalLines,
    filesCreated,
    totalExtractedLines,
    remainingLines,
    filesCreatedDetails,
    remainingContent,
    totalCodeblocks,
  };
}

module.exports = { processAIResponse }; 